# -*- coding: utf-8 -*-

import torch
import numpy as np
import torch.nn as nn

""" battle_v2 

    env.observation_spaces['blue_80'].shape -> (13, 13, 41)
    env.action['blue_80'].n -> 21
"""


class RNN(nn.Module):

    def __init__(self, agent_num, obs_space, action_space):
        super(RNN, self).__init__()
        """ Policy Network: CNN -> FC -> GRU -> FC 
            如何在RNN部分把不同的agent的区别给显示出来呢？
            既要让各个智能体共享大部分网络参数，又必须要表现出区别来。
        """

        # CNN
        self.conv1 = nn.Conv2d(in_channels=41, out_channels=82,
                               kernel_size=5, padding=2, stride=2)
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        self.conv2 = nn.Conv2d(in_channels=82, out_channels=164,
                               kernel_size=5, padding=1, stride=1)

        self.cnn_out_channel = 164
        self.obs_embedding_dim = 64

        # MLP
        self.mlp = nn.Sequential(nn.Linear(self.cnn_out_channel, self.obs_embedding_dim),
                                 nn.ReLU(True))

        self.rnn = nn.GRUCell(input_size=self.obs_embedding_dim,
                              hidden_size=self.obs_embedding_dim)

        # MLP
        self.output = nn.Linear(self.obs_embedding_dim, action_space)

    def forward(self, obs, hidden_state):
        """
            obs: (batch_size, 13, 13, 41)
            CNN -> MLP -> GRU -> MLP
        """

        obs = obs.reshape(-1, 13, 13, 41)
        obs = obs.permute(0, 3, 1, 2)

        # CNN
        e = self.conv1(obs)
        e = self.pool1(e)
        e = self.conv2(e)
        e = e.view(-1, self.cnn_out_channel)

        # MLP
        x = self.mlp(e)
        print(x.shape)

        # GRU
        # h = hidden_state.reshape(self.obs_embedding_dim)
        h = self.rnn(x, hidden_state)

        # MLP
        q = self.fc2(h)

        return q, h


if __name__ == '__main__':
    test_ts = np.random.randn(8, 13, 13, 41)
    test_ts = torch.from_numpy(test_ts).float()

    test_action = np.random.randn(64, 64)
    test_action = torch.from_numpy(test_action).float()

    model = RNN((13, 13, 41), 21)
    out = model(test_ts, test_action)
    print(out.shape)
